;; Copyright 2009 Adam C. Foltzer

;; Licensed under the Apache License, Version 2.0 (the "License");
;; you may not use this file except in compliance with the License.
;; You may obtain a copy of the License at

;; http://www.apache.org/licenses/LICENSE-2.0

;; Unless required by applicable law or agreed to in writing, software
;; distributed under the License is distributed on an "AS IS" BASIS,
;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
;; See the License for the specific language governing permissions and
;; limitations under the License.

(load "mmath.scm")

(define make-random-vector
  (lambda (n rand-min rand-max)
    (if (= rand-min rand-max)
        (make-vector n rand-min)
        (let* ([v (make-vector n)]
               [map-fn (lambda (x) (+ rand-min (random (- rand-max rand-min))))])
          (vector-map map-fn v)))))
  
;; Particles of the swarm are represented as vectors. The indices are tied to
;; certain values, although helpers exist to abstract away from this:
;; index 0 := vector representing position
;; index 1 := vector representing velocity
;; index 2 := vector with position of personal best
;; index 3 := fitness value of personal best
;; index 4 := neighborhood id

(define make-particle
  (lambda (pos vel best-val best-fit neigh)
    (vector pos vel best-val best-fit neigh)))

(define position
  (lambda (p)
    (vector-ref p 0)))

(define velocity
  (lambda (p)
    (vector-ref p 1)))

(define personal-best-pos
  (lambda (p)
    (vector-ref p 2)))

(define personal-best-val
  (lambda (p)
    (vector-ref p 3)))

(define neighborhood
  (lambda (p)
    (vector-ref p 4)))

;; Fitness functions are used throughout, usually denoted by f. An f
;; takes a position, and returns some fitness value for that positionA
;; useful f might be a function that returns the distance between the
;; particle and a known target point.

;; Initializes a population of neighborhoods, each with n paricles in
;; a random distribution across the specified space. Two vectors are
;; used to specify the dimensionality and size of the space: one
;; vector for minimum values of each dimension, and a corresponding
;; vector with maximum values. For example, (initialize-population 10
;; #(0.0 0.0) #(100.0 50.0 (lambda (p) (mm-distance (position p)
;; #(50.0, 25.0)))) creates a swarm with 10 particles across a
;; 100.0x50.0 space.  The personal best is initialized with the
;; starting position according to the fitness function f.

(define initialize-population
  (lambda (neighs n min-v max-v f)
    (let* ([population (make-vector (* neighs n))]
           [this-n 0]
           [dims (vector-length min-v)]
           [map-fn (lambda (x)
                     (let ([pos (make-vector dims)]
                           [vel (make-vector dims 0.0)])
                       (let loop ([index 0])
                         (if (< index dims)
                             (let* ([min (vector-ref min-v index)]
                                    [max (vector-ref max-v index)]
                                    [range (exact->inexact (- max min))])
                               (vector-set! pos index (+ (random range) min))
                               (loop (add1 index)))))
                       (set! this-n (modulo (add1 this-n) neighs))
                       (make-particle pos vel pos (f pos) this-n)))])
      (vector-map map-fn population))))

;; Returns the most-fit member of the population according to the
;; provided fitness values of the particles and optimization
;; comparator (usually > or <). Takes the number of neighborhoods.
;; Takes the output from fittest-neighbors as the input for fits.

(define fittest-member
  (lambda (fits comp neighs)
    (let ([v-length (vector-length fits)])
      (let loop ([index 0]
                 [best-val (personal-best-val (vector-ref fits 0))]
                 [best-index 0])
        (if (= index v-length)
            (vector-ref fits best-index)
            (let ([this-val (personal-best-val (vector-ref fits index))])
              (if (comp this-val best-val)
                  (loop (add1 index) this-val index)
                  (loop (add1 index) best-val best-index))))))))

(define fittest-neighbors
  (lambda (pop comp neighs)
    (let* ([v-length (vector-length pop)]
           [uninit (gensym)]
           [fits (make-vector neighs uninit)])
      (let loop ([index 0])
        (if (= index v-length)
            fits
            (let* ([this-p (vector-ref pop index)]
                   [this-val (personal-best-val this-p)]
                   [this-neigh (neighborhood this-p)]
                   [neigh-best (vector-ref fits this-neigh)])
              (if (or (eq? neigh-best uninit)
                      (comp this-val (personal-best-val neigh-best)))
                  (begin (vector-set! fits this-neigh this-p) (loop (add1 index)))
                  (loop (add1 index)))))))))

;; Updates fitness values for members by applying the fitness function
;; to the current position, updating the personal best according to
;; the optimization comparator

(define update-fitnesses
  (lambda (pop f comp)
    (let ([map-fn (lambda (p)
                    (let* ([pos (position p)]
                           [vel (velocity p)]
                           [best-pos (personal-best-pos p)]
                           [best-val (personal-best-val p)]
                           [neigh (neighborhood p)]
                           [this-val (f pos)])
                      (if (comp this-val best-val)
                          (make-particle pos vel pos this-val neigh)
                          (make-particle pos vel best-pos best-val neigh))))])
      (vector-map map-fn pop))))

;; Ticks the swarm once, updating fitnesses according to the provided
;; fitness function and optimization comparator. Returns the new population.

;; Simulation parameters:
;; neighs := number of neighborhoods
;; inertia := inertial constant slows velocity before tick adds new velocity
;; cognitive := weight toward personal best
;; neighw := weight toward neighborhood best
;; social := weight toward global best
;; r1-min/max := stochastic coefficient for cognitive weight
;; r2-min/max := stochastic coefficient for neighborhood weight
;; r3-min/max := stochastic coefficient for global weight

(define tick-swarm
  (lambda (pop f comp neighs inertia neighw cognitive social r1-min r1-max r2-min r2-max r3-min r3-max)
    (let* ([fits (fittest-neighbors pop comp neighs)]
           [global-best (fittest-member fits comp neighs)]
           [global-best-pos (personal-best-pos global-best)]
           [global-best-val (personal-best-val global-best)]
           [dims (vector-length global-best-pos)]
           [map-fn
            (lambda (p)
              (let* ([rand-v1 (make-random-vector dims r1-min r1-max)]
                     [rand-v2 (make-random-vector dims r2-min r2-max)]
                     [rand-v3 (make-random-vector dims r3-min r3-max)]
                     [best-neigh (vector-ref fits (neighborhood p))]
                     [new-vel
                      (mm-matrix-sum
                       (mm-scalar-multiply (velocity p) inertia)
                       (mm-scalar-multiply
                        (mm-entrywise-product (mm-matrix-diff (personal-best-pos p) (position p)) rand-v1)
                        cognitive)
                       (mm-scalar-multiply
                        (mm-entrywise-product (mm-matrix-diff (personal-best-pos best-neigh) (position p)) rand-v2)
                        neighw)
                       (mm-scalar-multiply
                        (mm-entrywise-product (mm-matrix-diff global-best-pos (position p)) rand-v3)
                        social))]
                     [new-pos
                      (mm-matrix-sum (position p) new-vel)])
                (make-particle new-pos new-vel (personal-best-pos p) (personal-best-val p) (neighborhood p))))]
           [new-pop (vector-map map-fn pop)])
      (update-fitnesses new-pop f comp))))

;;TODO: make this friendlier and check for early convergence

(define converged?
  (lambda (p)
    (let ([v-length (vector-length p)]
          [margin .5])
      (let loop ([index 0])
        (or (= index v-length)
            (and (< (personal-best-val (vector-ref p index)) margin)
                 (loop (add1 index))))))))

(define run-swarm
  (lambda (pop f comp neighs inertia cognitive neighw social r1-min r1-max r2-min r2-max r3-min r3-max iter)
    (let loop ([i 0] [pop pop])
      (cond
        [(converged? pop) (printf "~f,~f,~f,~s\n" neighw social r2-min i)]
        [(< i iter)
         (loop (add1 i)
               (tick-swarm pop f comp neighs inertia cognitive neighw social r1-min r1-max r2-min r2-max r3-min r3-max))]
        [else (printf "~f,~f,~f,~s\n" neighw social r2-min i)]))))

(define go
  (lambda ()
    (let* ([min-v #(0 0 0)]
           [max-v #(100 100 100)]
           [target #(50 50 50)]
           [fit-fn (lambda (p) (mm-distance p target))]
           [comp <]
           [neighs 5]
           [pop-per-neigh 20]
           [pop (initialize-population neighs pop-per-neigh min-v max-v fit-fn)]
           [inertia 0.8]
           [cognitive 0.8]
           [neighw 0.8]
           [social 0.2]
           [r1-min 0]
           [r1-max 1.0]
           [r2-min 0.0]
           [r2-max 1.0]
           [r3-min 0.0]
           [r3-max 1.0]
           [iters 1500000])

      (run-swarm pop fit-fn comp neighs inertia cognitive neighw social r1-min r1-max r2-min r2-max r3-min r3-max iters))))

'(pretty-print (go))

(define map-go (lambda (foo) (go)))